import time
from typing import Any

import cv2
import numpy
from insightface.app import FaceAnalysis
from insightface.model_zoo.inswapper import INSwapper

from Settings import ModelsBaseDIR


class ChangeFace:
    def __init__(self):
        self.__FaceAnalysis = self.__initFaceAnalysisEngine()
        self.__FaceSwapper = self.__initFaceSwapperEngine()

    def __initFaceAnalysisEngine(self):
        face_analysis = FaceAnalysis(
            # 下载模型的目录名称
            name='buffalo_l',
            # 模型下载位置
            root=ModelsBaseDIR,
            # 'AzureExecutionProvider, CPUExecutionProvider'
            providers=["CPUExecutionProvider"]
        )
        # 预处理
        face_analysis.prepare(ctx_id=0, det_size=(640, 640))

        return face_analysis

    def __initFaceSwapperEngine(self):
        inswapper = INSwapper(
            model_file=fr'C:\Users\Administrator\Desktop\change_face\change_face\models\1.onnx',
        )

        return inswapper

    def _getOneFace(self, img_path):
        frame = cv2.imread(img_path)
        face = self.__FaceAnalysis.get(frame)
        try:
            # 如果拿到了多张脸，就只返回最近的一个，返回了一个Face的对象
            return min(face, key=lambda x: x.bbox[0])
        except ValueError:
            raise "传入图片地址"

    def _swapFace(self, tmp_face, target_face, source_face) -> numpy.ndarray[Any, Any]:
        frame = self.__FaceSwapper.get(tmp_face, target_face, source_face, paste_back=True)
        return frame

    def start(self,target_path, source_path, save__path):
        source_face = self._getOneFace(source_path)
        target_face = self._getOneFace(target_path)

        tmp_region_result: numpy.ndarray[Any, Any] = cv2.imread(target_path)

        frame = self._swapFace(tmp_region_result, target_face, source_face)
        cv2.imwrite(save__path, frame)

        pass

if __name__ == '__main__':
    start_time = time.time()

    changeFace = ChangeFace()

    from Settings import BaseTestDataCubePath, OutBaseDir
    source_image_path = f'{BaseTestDataCubePath}\\1.jpg'
    target_image_path = f'{BaseTestDataCubePath}\\4.jpg'

    save_image_path = f'{OutBaseDir}\\1.jpg'
    save_image_path1 = f'{OutBaseDir}\\2.jpg'

    changeFace.start(source_image_path, target_image_path, save_image_path)
    print(time.time() - start_time)

    start_time1 = time.time()
    changeFace.start(source_image_path, target_image_path, save_image_path1)
    changeFace.start(source_image_path, target_image_path, save_image_path1)
    print(time.time() - start_time1)